{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our task is to build a consultant bot that can answer questions of different domains, such as medical with a doctor bot or legal with a lawyer bot.\n",
    "\n",
    "We will show how flexible ``Component`` and the ``Sequential`` container is to build the same task\n",
    "in different ways.\n",
    "\n",
    "1. **Single Task**: We can build a single task where it deals with multiple generators and handles any coding logic.\n",
    "2. **Multiple Tasks** and combine them using ``Sequential`` which resembles the concept of `Chain` or pipelines in other libraries."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, lets prepare the imports and prompt templates using `jinjia2` template. We plan to demonstrate how we can use different models too. If this tutorial is the first thing you read, no need to care about more details, but focus on how the `development process` looks like using `LightRAG` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from lightrag.core import Component, Generator, Sequential\n",
    "from lightrag.components.model_client import OpenAIClient\n",
    "from lightrag.components.model_client import GroqAPIClient\n",
    "from lightrag.utils import setup_env # make sure you have a .env file with OPENAI_API_KEY and GROQ_API_KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "template_doc = r\"\"\"<SYS> You are a doctor </SYS> User: {{input_str}}\"\"\"\n",
    "template_law = r\"\"\"<SYS> You are a lawyer </SYS> User: {{input_str}}\"\"\"\n",
    "template_router = r\"\"\"<SYS> You are a router who will route a user question to the right generator.\n",
    "            Here are your choices in form of key: value pairs:\n",
    "             {% for key, value in choices.items() %}\n",
    "                {{ key }}: {{ value }}\n",
    "             {% endfor %}\n",
    "            Output the key of your choice.\n",
    "            </SYS> User question: {{input_str}}\n",
    "            You:\n",
    "            \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's turn on the library log to help with debugging.\n",
    "from lightrag.utils import enable_library_logging\n",
    "enable_library_logging()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is our first approach to build a single task with multiple generators and call each conditionally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChatBotWithRouter(Component):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        model_1_kwargs = {\n",
    "            \"model\": \"gpt-3.5-turbo\",\n",
    "        }\n",
    "        model_2_kwargs = {\"model\": \"llama3-8b-8192\"}\n",
    "        self.doc = Generator(\n",
    "            template=template_doc,\n",
    "            model_client=OpenAIClient(),\n",
    "            model_kwargs=model_1_kwargs,\n",
    "        )\n",
    "        self.lawyer = Generator(\n",
    "            template=template_law,\n",
    "            model_client=GroqAPIClient(),\n",
    "            model_kwargs=model_2_kwargs,\n",
    "        )\n",
    "        self.router_choices = {\n",
    "            \"doctor\": self.create_generator_signature(self.doc),\n",
    "            \"lawyer\": self.create_generator_signature(self.lawyer),\n",
    "            \"other\": \"Choose me the question does not apply to other choices.\",\n",
    "        }\n",
    "        print(self.router_choices)\n",
    "\n",
    "        self.router = Generator(\n",
    "            template=template_router,\n",
    "            model_client=OpenAIClient(),\n",
    "            model_kwargs=model_1_kwargs,\n",
    "        )\n",
    "\n",
    "    def call(self, query: str) -> str:\n",
    "        choice = self.router(\n",
    "            prompt_kwargs={\"input_str\": query, \"choices\": self.router_choices}\n",
    "        ).data\n",
    "        if choice == \"doctor\":\n",
    "            return self.doc(prompt_kwargs={\"input_str\": query}).data\n",
    "        elif choice == \"lawyer\":\n",
    "            return self.lawyer(prompt_kwargs={\"input_str\": query}).data\n",
    "        else:\n",
    "            return \"Sorry, I cannot help you with that.\"\n",
    "\n",
    "    def create_generator_signature(self, generator: Generator):\n",
    "        template = generator.template\n",
    "        pattern = r\"<SYS>(.*?)</SYS>\"\n",
    "\n",
    "        matches = re.findall(pattern, template)\n",
    "        for match in matches:\n",
    "            print(\"Content between <SYS> tags:\", match)\n",
    "            return match"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 23:13:16 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str']\n",
      "2024-06-09 23:13:16 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str']\n",
      "Content between <SYS> tags:  You are a doctor \n",
      "Content between <SYS> tags:  You are a lawyer \n",
      "{'doctor': ' You are a doctor ', 'lawyer': ' You are a lawyer ', 'other': 'Choose me the question does not apply to other choices.'}\n",
      "2024-06-09 23:13:16 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str', 'choices']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ChatBotWithRouter(\n",
       "  (doc): Generator(\n",
       "    model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
       "    (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "    (model_client): OpenAIClient()\n",
       "  )\n",
       "  (lawyer): Generator(\n",
       "    model_kwargs={'model': 'llama3-8b-8192'}, model_type=ModelType.LLM\n",
       "    (prompt): Prompt(template: <SYS> You are a lawyer </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "    (model_client): GroqAPIClient()\n",
       "  )\n",
       "  (router): Generator(\n",
       "    model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
       "    (prompt): Prompt(\n",
       "      template: <SYS> You are a router who will route a user question to the right generator.\n",
       "                  Here are your choices in form of key: value pairs:\n",
       "                   {% for key, value in choices.items() %}\n",
       "                      {{ key }}: {{ value }}\n",
       "                   {% endfor %}\n",
       "                  Output the key of your choice.\n",
       "                  </SYS> User question: {{input_str}}\n",
       "                  You:\n",
       "                  , prompt_variables: ['input_str', 'choices']\n",
       "    )\n",
       "    (model_client): OpenAIClient()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initiate the task component, and print the task details.\n",
    "\n",
    "task = ChatBotWithRouter()\n",
    "task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 23:13:16 - INFO - [generator.py:194:call] - prompt_kwargs: {'input_str': 'I have a legal question', 'choices': {'doctor': ' You are a doctor ', 'lawyer': ' You are a lawyer ', 'other': 'Choose me the question does not apply to other choices.'}}\n",
      "2024-06-09 23:13:16 - INFO - [generator.py:195:call] - model_kwargs: {}\n",
      "2024-06-09 23:13:16 - INFO - [openai_client.py:122:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<SYS> You are a router who will route a user question to the right generator.\\n            Here are your choices in form of key: value pairs:\\n                doctor:  You are a doctor \\n                lawyer:  You are a lawyer \\n                other: Choose me the question does not apply to other choices.\\n            Output the key of your choice.\\n            </SYS> User question: I have a legal question\\n            You:'}]}\n",
      "2024-06-09 23:13:17 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 23:13:17 - INFO - [generator.py:203:call] - output: GeneratorOutput(data='lawyer', error=None, raw_response='lawyer')\n",
      "2024-06-09 23:13:17 - INFO - [generator.py:194:call] - prompt_kwargs: {'input_str': 'I have a legal question'}\n",
      "2024-06-09 23:13:17 - INFO - [generator.py:195:call] - model_kwargs: {}\n",
      "2024-06-09 23:13:17 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 23:13:17 - INFO - [generator.py:203:call] - output: GeneratorOutput(data=\"Please go ahead and share your legal question with me. I'll do my best to provide you with a clear and concise answer based on my knowledge of the law.\", error=None, raw_response=\"Please go ahead and share your legal question with me. I'll do my best to provide you with a clear and concise answer based on my knowledge of the law.\")\n",
      "Please go ahead and share your legal question with me. I'll do my best to provide you with a clear and concise answer based on my knowledge of the law.\n"
     ]
    }
   ],
   "source": [
    "# Call the task with a query\n",
    "\n",
    "query = \"I have a legal question\"\n",
    "print(task(query))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's separate this into multiple subtasks and ``chain`` them using the ``Sequential`` container.\n",
    "\n",
    "First, the router task which will takes a dictionary of choices and return the selected key. In addition, we use ``_extra_repr`` to improve the default string representation of the task.\n",
    "\n",
    "As ``Sequential`` will pass the output of one task to the next using positional arguments, we return whatever is needed to the next task in a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Router component\n",
    "\n",
    "from typing import Dict\n",
    "class Router(Component):\n",
    "    def __init__(self, choices: Dict[str, str] = {}):\n",
    "        super().__init__()\n",
    "        self.choices = choices\n",
    "        self.router = Generator(\n",
    "            template=template_router,\n",
    "            model_client=OpenAIClient(),\n",
    "            model_kwargs={\"model\": \"gpt-3.5-turbo\"},\n",
    "        )\n",
    "\n",
    "    def call(self, query: str) -> str:\n",
    "        prompt_kwargs = {\"input_str\": query, \"choices\": self.choices}\n",
    "        choice =  self.router(prompt_kwargs=prompt_kwargs).data\n",
    "        return {\"choice\": choice, \"query\": query}\n",
    "    \n",
    "    def _extra_repr(self):\n",
    "        return f\"Choices: {self.choices}, \""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 23:13:17 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str', 'choices']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Router(\n",
       "  Choices: {}, \n",
       "  (router): Generator(\n",
       "    model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
       "    (prompt): Prompt(\n",
       "      template: <SYS> You are a router who will route a user question to the right generator.\n",
       "                  Here are your choices in form of key: value pairs:\n",
       "                   {% for key, value in choices.items() %}\n",
       "                      {{ key }}: {{ value }}\n",
       "                   {% endfor %}\n",
       "                  Output the key of your choice.\n",
       "                  </SYS> User question: {{input_str}}\n",
       "                  You:\n",
       "                  , prompt_variables: ['input_str', 'choices']\n",
       "    )\n",
       "    (model_client): OpenAIClient()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r = Router()\n",
    "r"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets build another subtask which handles the chat depending on the selected key from the router task.\n",
    "As the router task returns a dictionary, we will make our input dictionary type that parses the ``choice`` and ``query`` key value pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the second chat component with two generators\n",
    "\n",
    "class Chat(Component):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.doc = Generator(\n",
    "            template=template_doc,\n",
    "            model_client=OpenAIClient(),\n",
    "            model_kwargs={\"model\": \"gpt-3.5-turbo\"},\n",
    "        )\n",
    "        self.lawyer = Generator(\n",
    "            template=template_law,\n",
    "            model_client=GroqAPIClient(),\n",
    "            model_kwargs={\"model\": \"llama3-8b-8192\"},\n",
    "        )\n",
    "    # to chain together just to make sure the output can be directly passed to the next as input\n",
    "    def call(self, input: Dict[str, str]) -> Dict[str, str]:\n",
    "        choice = input.get(\"choice\", None)\n",
    "        query = input.get(\"query\", None)\n",
    "        if choice == \"doctor\":\n",
    "            return self.doc(prompt_kwargs={\"input_str\": query}).data\n",
    "        elif choice == \"lawyer\":\n",
    "            return self.lawyer(prompt_kwargs={\"input_str\": query}).data\n",
    "        else:\n",
    "            return \"Sorry, I am not able to help you with that.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 23:13:17 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str']\n",
      "2024-06-09 23:13:17 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Chat(\n",
       "  (doc): Generator(\n",
       "    model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
       "    (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "    (model_client): OpenAIClient()\n",
       "  )\n",
       "  (lawyer): Generator(\n",
       "    model_kwargs={'model': 'llama3-8b-8192'}, model_type=ModelType.LLM\n",
       "    (prompt): Prompt(template: <SYS> You are a lawyer </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "    (model_client): GroqAPIClient()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat = Chat()\n",
    "chat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets chain the router and the chat task using the ``Sequential`` container into a runnable pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QAWithRouter(Component):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.router = Router(choices={\"doctor\": \"Doctor\", \"lawyer\": \"Lawyer\", \"other\": \"Other\"})\n",
    "        self.chat = Chat()\n",
    "        self.pipeline = Sequential(self.router, self.chat)\n",
    "\n",
    "    def call(self, query: str) -> str:\n",
    "        return self.pipeline(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 23:13:17 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str', 'choices']\n",
      "2024-06-09 23:13:17 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str']\n",
      "2024-06-09 23:13:17 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "QAWithRouter(\n",
       "  (router): Router(\n",
       "    Choices: {'doctor': 'Doctor', 'lawyer': 'Lawyer', 'other': 'Other'}, \n",
       "    (router): Generator(\n",
       "      model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
       "      (prompt): Prompt(\n",
       "        template: <SYS> You are a router who will route a user question to the right generator.\n",
       "                    Here are your choices in form of key: value pairs:\n",
       "                     {% for key, value in choices.items() %}\n",
       "                        {{ key }}: {{ value }}\n",
       "                     {% endfor %}\n",
       "                    Output the key of your choice.\n",
       "                    </SYS> User question: {{input_str}}\n",
       "                    You:\n",
       "                    , prompt_variables: ['input_str', 'choices']\n",
       "      )\n",
       "      (model_client): OpenAIClient()\n",
       "    )\n",
       "  )\n",
       "  (chat): Chat(\n",
       "    (doc): Generator(\n",
       "      model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
       "      (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "      (model_client): OpenAIClient()\n",
       "    )\n",
       "    (lawyer): Generator(\n",
       "      model_kwargs={'model': 'llama3-8b-8192'}, model_type=ModelType.LLM\n",
       "      (prompt): Prompt(template: <SYS> You are a lawyer </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "      (model_client): GroqAPIClient()\n",
       "    )\n",
       "  )\n",
       "  (pipeline): Sequential(\n",
       "    (0): Router(\n",
       "      Choices: {'doctor': 'Doctor', 'lawyer': 'Lawyer', 'other': 'Other'}, \n",
       "      (router): Generator(\n",
       "        model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
       "        (prompt): Prompt(\n",
       "          template: <SYS> You are a router who will route a user question to the right generator.\n",
       "                      Here are your choices in form of key: value pairs:\n",
       "                       {% for key, value in choices.items() %}\n",
       "                          {{ key }}: {{ value }}\n",
       "                       {% endfor %}\n",
       "                      Output the key of your choice.\n",
       "                      </SYS> User question: {{input_str}}\n",
       "                      You:\n",
       "                      , prompt_variables: ['input_str', 'choices']\n",
       "        )\n",
       "        (model_client): OpenAIClient()\n",
       "      )\n",
       "    )\n",
       "    (1): Chat(\n",
       "      (doc): Generator(\n",
       "        model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
       "        (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "        (model_client): OpenAIClient()\n",
       "      )\n",
       "      (lawyer): Generator(\n",
       "        model_kwargs={'model': 'llama3-8b-8192'}, model_type=ModelType.LLM\n",
       "        (prompt): Prompt(template: <SYS> You are a lawyer </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "        (model_client): GroqAPIClient()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa = QAWithRouter()\n",
    "qa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 23:13:18 - INFO - [generator.py:194:call] - prompt_kwargs: {'input_str': 'I have a legal question', 'choices': {'doctor': 'Doctor', 'lawyer': 'Lawyer', 'other': 'Other'}}\n",
      "2024-06-09 23:13:18 - INFO - [generator.py:195:call] - model_kwargs: {}\n",
      "2024-06-09 23:13:18 - INFO - [openai_client.py:122:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<SYS> You are a router who will route a user question to the right generator.\\n            Here are your choices in form of key: value pairs:\\n                doctor: Doctor\\n                lawyer: Lawyer\\n                other: Other\\n            Output the key of your choice.\\n            </SYS> User question: I have a legal question\\n            You:'}]}\n",
      "2024-06-09 23:13:18 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 23:13:18 - INFO - [generator.py:203:call] - output: GeneratorOutput(data='lawyer', error=None, raw_response='lawyer')\n",
      "2024-06-09 23:13:18 - INFO - [generator.py:194:call] - prompt_kwargs: {'input_str': 'I have a legal question'}\n",
      "2024-06-09 23:13:18 - INFO - [generator.py:195:call] - model_kwargs: {}\n",
      "2024-06-09 23:13:18 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 23:13:18 - INFO - [generator.py:203:call] - output: GeneratorOutput(data=\"I'd be happy to help you with your legal question! Can you please provide some more details about your question, such as what type of law you're dealing with (e.g. criminal, family, employment, etc.) and what specific issues or concerns you're facing?\", error=None, raw_response=\"I'd be happy to help you with your legal question! Can you please provide some more details about your question, such as what type of law you're dealing with (e.g. criminal, family, employment, etc.) and what specific issues or concerns you're facing?\")\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"I'd be happy to help you with your legal question! Can you please provide some more details about your question, such as what type of law you're dealing with (e.g. criminal, family, employment, etc.) and what specific issues or concerns you're facing?\""
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa(\"I have a legal question\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my-project-kernel",
   "language": "python",
   "name": "my-project-kernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
